package main

import (
	"fmt"
	"time"
	"rand"
	"bytes"
	"os"
	"math"
	"sort"
	"flag"
)

// Representa um segundo em nanosegundos
const SECOND = 1e9

// Resposta do processamento
type Response int64

// Uma requisição enviada pelo cliente para ser processada
type Request struct {
	Id       int           // identificador da requisicao do cliente
	Pos      int           // identificador de quando chegou
	Response chan Response // canal para enviar o resultado
}

// Status do trabalhador
type Status int

// Trabalhador responsável por processar as requisições
type Worker struct {
	Id      string
	pending int           // tamanho da fila de trabalho do mesmo
	Request chan *Request // canal para receber as requisições
}

type WorkerPool []Worker

type sorter struct {
	bl *Balancer
}

// Funcao principal do trabalhador
func (w *Worker) Start(done chan Worker) {
	for req := range w.Request {
		// simula que esta processando algo
		ret := rand.Int63n(SECOND*int64(req.Id)) % (SECOND * 2)
		time.Sleep(ret)

		// efetua o cast e envia a resposta para o cliente
		req.Response <- Response(ret)
		done <- *w

		if *verbose {
			fmt.Printf("%7v\t%v\t-1\n", req.Pos, w.Id)
		}
	}
}

func (wp WorkerPool) String() string {
	ret := ""
	for _, el := range wp {
		ret = ret + fmt.Sprintf(" (%v)", el.pending)
	}
	return ret
}

func (s sorter) Less(i, j int) bool {
	wp := s.bl.wp
	return s.bl.load[wp[i].Id] < s.bl.load[wp[j].Id]
}

func (s sorter) Swap(i, j int) {
	wp := s.bl.wp
	aux := wp[i]
	wp[i] = wp[j]
	wp[j] = aux
}

func (s sorter) Len() int {
	return len(s.bl.wp)
}

func (s sorter) Next() (w Worker) {
	if *rrobin {
		w = s.bl.wp[s.bl.reqC%len(s.bl.wp)]
	} else {
		w = s.bl.wp[0]
	}
	return
}

func (s sorter) Update() {
	sort.Sort(s)
}

type Balancer struct {
	wp        WorkerPool // lista de trabalhadores
	load      map[string]int
	requests  chan *Request // canal para receber serviços
	done      chan Worker   // canal para notificar que o serviço está pronto
	reqC      int           // total de requisições
	algorithm sorter        // algoritmo de distribuição de carga
	exit      chan int      // indica o status de término do balancer
}

func NewBalancer(sz, factor int) (b *Balancer) {
	b = &Balancer{make(WorkerPool, sz), make(map[string]int), make(chan *Request, sz), make(chan Worker), 0,
		sorter{nil}, make(chan int)}
	b.algorithm.bl = b

	for i := 0; i < sz; i++ {
		b.wp[i] = *b.NewWorker(i)
		b.load[b.wp[i].Id] = 0
	}
	return
}

func (b *Balancer) dispatch(req *Request) {
	b.reqC++
	req.Pos = b.reqC
	w := b.algorithm.Next()
	w.Request <- req
	b.load[w.Id]++
	b.algorithm.Update()

	if *verbose {
		fmt.Printf("%7v\t%v\t1\n", b.reqC, w.Id)
	}
}

func (b *Balancer) receive(w Worker) {
	b.load[w.Id]--
	b.algorithm.Update()
}

func (b *Balancer) NewWorker(id int) (w *Worker) {
	w = &Worker{fmt.Sprintf("%3x", id), 0, make(chan *Request, 200)}
	go w.Start(b.done)
	return
}

func (b *Balancer) LoadAvg() float64 {
	sum := float64(0)
	for _, v := range b.load {
		sum += float64(v)
	}
	return sum / float64(len(b.load))
}

func (b *Balancer) DesvioPadrao() (avg float64, desvp float64) {
	avg = b.LoadAvg()
	aux := float64(0)
	for _, v := range b.load {
		aux += (float64(v) - avg) * (float64(v) - avg)
	}
	desvp = math.Sqrt(aux / float64(len(b.load)-1))
	return
}

func (b *Balancer) String() string {
	buff := bytes.NewBuffer(make([]byte, 0))
	fmt.Fprintf(buff, "r=%v/", b.reqC)
	fmt.Fprintf(buff, "w=%v/", len(b.wp))
	if len(b.wp) <= 10 {
		for _, v := range b.load {
			fmt.Fprintf(buff, " (%v)", v)
		}
	}
	avg, desvp := b.DesvioPadrao()
	fmt.Fprintf(buff, "/avg=%v/desvp=%v", avg, desvp)
	return buff.String()
}

func (b *Balancer) DistributeLoad() {
	t := time.NewTicker(SECOND / 2)
	for {
		select {
		case req := <-b.requests:
			b.dispatch(req)
		case w := <-b.done:
			b.receive(w)
			//                fmt.Printf("%v-> %s\n", b.reqC, b.wp)
		case <-t.C:
			fmt.Fprintf(os.Stdout, "%v\n", b)
		}
	}
}

func producer(req chan *Request) {
	i := int(0)
	for {
		i++
		time.Sleep(SECOND + rand.Int63n(SECOND))
		ret := make(chan Response)
		r := &Request{i, -1, ret}
		req <- r
		<-ret
	}
}

func producers(req chan *Request, pCount int) {
	t := time.NewTicker(SECOND / 100)
	for i := 0; i < pCount; i++ {
		go producer(req)
		<-t.C
	}
}

var (
	poolSize = flag.Int("w", 10, "Number of workers")
	factor   = flag.Int("f", 50, "Number of clients por worker")
	nclients = flag.Int("n", 500, "Number of clients")
	rrobin   = flag.Bool("rr", false, "Usa Round-Robin ao invés de balanceamento")
	verbose  = flag.Bool("v", false, "Verbose")
)

func main() {
	flag.Parse()
	bl := NewBalancer(*poolSize, *factor)
	go bl.DistributeLoad()
	go producers(bl.requests, *nclients)
	<-bl.exit
}
